package utility

import (
	"sort"

	"github.com/gogf/gf/v2/util/gconv"
)

// TreeNode 泛型树节点接口
type TreeNode[T any] interface {
	GetId() string
	GetParentId() string
	GetName() string
	GetSort() string
	GetChildren() []T
	SetChildren([]T)
}

// TreeTool 泛型树形工具
type TreeTool[T TreeNode[T]] struct{}

// ListToTree 列表转树(核心方法)
func (t *TreeTool[T]) ListToTree(list []T, parentId string) []T {
	if len(list) == 0 {
		return []T{}
	}

	// 使用map存储父子关系
	treeMap := make(map[string][]T)
	for _, item := range list {
		pid := item.GetParentId()
		treeMap[pid] = append(treeMap[pid], item)
	}

	// 递归构建树
	var buildTree func(string) []T
	buildTree = func(id string) []T {
		nodes := treeMap[id]
		if len(nodes) == 0 {
			return nil
		}

		// 排序
		sort.Slice(nodes, func(i, j int) bool {
			sort1 := gconv.Int64(nodes[i].GetSort())
			sort2 := gconv.Int64(nodes[j].GetSort())
			return sort1 < sort2
		})

		// 设置子节点
		for i := range nodes {
			children := buildTree(nodes[i].GetId())
			nodes[i].SetChildren(children)
		}

		return nodes
	}

	return buildTree(parentId)
}

// Recursion 递归遍历树
func (t *TreeTool[T]) Recursion(
	treeData []T,
	beforeFunc func(T, interface{}),
	afterFunc func(T, interface{}),
	res interface{},
) {
	if len(treeData) == 0 {
		return
	}

	for _, data := range treeData {
		if beforeFunc != nil {
			beforeFunc(data, res)
		}

		t.Recursion(data.GetChildren(), beforeFunc, afterFunc, res)

		if afterFunc != nil {
			afterFunc(data, res)
		}
	}
}

// FindNode 查找指定节点
func (t *TreeTool[T]) FindNode(treeData []T, id string) (T, bool) {
	for _, node := range treeData {
		if node.GetId() == id {
			return node, true
		}
		if children := node.GetChildren(); len(children) > 0 {
			if found, ok := t.FindNode(children, id); ok {
				return found, true
			}
		}
	}
	var zero T
	return zero, false
}

// Flatten 将树展平为列表
func (t *TreeTool[T]) Flatten(treeData []T) []T {
	var result []T
	t.Recursion(treeData,
		func(node T, _ interface{}) {
			result = append(result, node)
		},
		nil,
		nil,
	)
	return result
}
